Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

@sanepunk thanks for the PR. I'll ask the team whether we would like to have a reinforcement learning example in flax repository (the maintenance of this code etc) and let you know here about the decision.

Thanks, @vfdev-5! I noticed there wasn’t an NNX-based RL example in the repo, so I thought it might be helpful to contribute one. Looking forward to hearing the team’s decision.

@sanepunk unfortunately, for now it was decided not to add new RL example to flax examples. This is mainly due to multiple reasons like:

• we should refresh and update existing examples such they use the latest Flax API (e.g. >=0.12.0)
• your code maintenance costs in future: keep versions up to date and fixing all BC-breaking changes. For example, previously, the change from gym to gymnasium and BC-breaking changes were a bit painful to dig into.
• discussable dependency: gymnax which has its last commit 5 months ago and whether we can rely on it.

Again thanks a lot for this RL example with Flax! I understand that this decision can be disappointing as you already done a decent amount of work, but please consider a good practice is to open first an issue in the repository about what people would like to add as contribution and get a feedback from maintainers.

@vfdev-5 Thank you so much for taking the time to review and explain this! I really appreciate the feedback and will definitely open an issue first next time.

Would it be okay if I open an issue to discuss adding NNX implementations to the existing nn module examples, since it could help people understand and learn the NNX module better?

Would it be okay if I open an issue to discuss adding NNX implementations to the existing nn module examples, since it could help people understand and learn the NNX module better?

Please open first a new issue to discuss first what would you like to do. I would say the priority is to make refresh the documentation, guides etc.

Let me close this PR following the reason given in: #5036 (comment)